# developer: Qingtian Liu.
# Date: 2024-10-7

import os
import re
import difflib
from tqdm import tqdm
from datetime import datetime

# 计算两个代码块的相似度
def compute_similarity(code1, code2):
    sm = difflib.SequenceMatcher(None, code1, code2)
    return sm.ratio()

# 获取文件夹中所有 .py 文件的数量
def count_py_files(root_directory):
    py_file_count = 0
    for dirpath, _, filenames in os.walk(root_directory):
        for filename in filenames:
            if filename.endswith(".py"):
                py_file_count += 1
    return py_file_count

# 统计并过滤重复代码块，使用相似度阈值
def count_lines_in_py_files(root_directory, similarity_threshold=0.92):
    file_stats = []
    processed_blocks = []  # 存储已处理过的代码块
    duplicate_blocks = []  # 存储重复代码块信息

    total_files = 0
    total_code_lines = 0
    total_empty_lines = 0
    total_comment_lines = 0
    total_docstring_lines = 0
    total_lines = 0

    # 获取文件总数，用于进度条显示
    total_py_files = count_py_files(root_directory)

    # 初始化进度条
    with tqdm(total=total_py_files, desc="Scanning Python files", unit="file") as pbar:
        for dirpath, _, filenames in os.walk(root_directory):
            for filename in filenames:
                if filename.endswith(".py"):
                    total_files += 1
                    py_file = os.path.join(dirpath, filename)
                    with open(py_file, 'r', encoding='utf-8') as f:
                        code_lines = 0
                        empty_lines = 0
                        comment_lines = 0
                        docstring_lines = 0
                        in_docstring = False
                        code_block = ""
                        blocks = []  # 用于存储代码块

                        for line in f:
                            stripped_line = line.strip()
                            total_lines += 1

                            # 空行
                            if not stripped_line:
                                empty_lines += 1
                            # 单行注释
                            elif stripped_line.startswith("#"):
                                comment_lines += 1
                            # 文档字符串 (开始或结束)
                            elif (stripped_line.startswith('"""') or stripped_line.startswith("'''")):
                                if in_docstring:
                                    docstring_lines += 1
                                    in_docstring = False
                                else:
                                    docstring_lines += 1
                                    in_docstring = True
                            # 文档字符串的内部行
                            elif in_docstring:
                                docstring_lines += 1
                            # 函数定义或类定义，代表代码块的结束
                            elif stripped_line.startswith("def ") or stripped_line.startswith("class "):
                                if code_block.strip():  # 遇到新的函数或类定义时，保存之前的代码块
                                    blocks.append(code_block)
                                code_block = stripped_line
                            else:
                                code_block += "\n" + stripped_line

                        # 处理最后的代码块
                        if code_block.strip():
                            blocks.append(code_block)

                        # 统计非重复代码块
                        unique_code_lines = 0
                        for block in blocks:
                            is_duplicate = False
                            for processed_block, file_info in processed_blocks:
                                similarity = compute_similarity(block, processed_block)
                                if similarity >= similarity_threshold:
                                    is_duplicate = True
                                    # 将换行符替换为空格 * 2，避免在 Markdown 中破坏表格布局
                                    formatted_block = block.replace("\n", "  ")
                                    formatted_processed_block = processed_block.replace("\n", "  ")
                                    duplicate_blocks.append({
                                        "filename": py_file,
                                        "similar_file": file_info["filename"],
                                        "duplicate_block": formatted_block[:50] + "...",
                                        "similar_block": formatted_processed_block[:50] + "...",
                                        "similarity": f"**{similarity:.2f}**",  # 高亮显示相似度
                                    })
                                    break
                            if not is_duplicate:
                                processed_blocks.append((block, {"filename": py_file}))
                                unique_code_lines += len(block.splitlines())

                        # 更新统计
                        total_code_lines += unique_code_lines

                        file_stats.append({
                            "filename": py_file,
                            "code_lines": unique_code_lines,
                            "empty_lines": empty_lines,
                            "comment_lines": comment_lines,
                            "docstring_lines": docstring_lines,
                            "total_lines": unique_code_lines + empty_lines + comment_lines + docstring_lines
                        })

                    # 每处理完一个文件，更新进度条
                    pbar.update(1)

    # 按相似度从高到低排序
    duplicate_blocks.sort(key=lambda x: x["similarity"], reverse=True)

    # 按代码行数降序排序
    file_stats.sort(key=lambda x: x["code_lines"], reverse=True)
    summary = {
        "total_files": total_files,
        "total_code_lines": total_code_lines,
        "total_empty_lines": total_empty_lines,
        "total_comment_lines": total_comment_lines,
        "total_docstring_lines": total_docstring_lines,
        "total_lines": total_lines
    }

    return file_stats, summary, duplicate_blocks

# 保存 Markdown 文件
def save_results_to_md(stats, summary, duplicate_blocks, root_directory):
    date_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    md_filename = os.path.join(root_directory, f"python_code_stats_{date_str}.md")
    
    with open(md_filename, "w", encoding="utf-8") as f:
        f.write(f"# Summary\n\n")
        f.write(f"Date : {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
        f.write(f"Directory: {root_directory}\n\n")
        f.write(f"Total: {summary['total_files']} files, {summary['total_code_lines']} unique code lines, "
                f"{summary['total_comment_lines']} comments, {summary['total_empty_lines']} blanks, "
                f"all {summary['total_lines']} lines (including duplicates)\n\n")

        # 文件统计细节
        f.write(f"## Files\n")
        f.write("| Filename | Unique Code Lines | Empty Lines | Comment Lines | Docstring Lines | Total Lines |\n")
        f.write("|---|---:|---:|---:|---:|---:|\n")
        for stat in stats:
            f.write(f"| {stat['filename']} | {stat['code_lines']} | {stat['empty_lines']} | "
                    f"{stat['comment_lines']} | {stat['docstring_lines']} | {stat['total_lines']} |\n")

        # 重复代码信息
        if duplicate_blocks:
            f.write(f"\n## Duplicate Code Blocks (Similarity >= 92%)\n\n")
            f.write("| Files (filename and similar file) | Origin Code Block | Similar Code Block | Similarity |\n")
            f.write("|---|---|---|---:|\n")
            for dup in duplicate_blocks:
                f.write(f"| {dup['filename']} {dup['similar_file']} | `{dup['duplicate_block']}` | "
                        f"`{dup['similar_block']}` | {dup['similarity']} |\n")

    print(f"Markdown results saved to {md_filename}")
    return md_filename

# 主函数：保存 Markdown
def main():
    root_directory = input("Please enter the directory to scan: ")
    if os.path.isdir(root_directory):
        stats, summary, duplicate_blocks = count_lines_in_py_files(root_directory)
        
        # 保存 Markdown 文件
        md_filename = save_results_to_md(stats, summary, duplicate_blocks, root_directory)
    else:
        print("The directory does not exist.")

if __name__ == "__main__":
    main()
